absl-py==0.11.0
chex==0.0.4
clu==0.0.2
einops==0.3.0
flax==0.3.2
jax==0.2.9
matplotlib==3.3.4
ml-collections==0.1.0
numpy==1.20.1
optax==0.0.2
scipy==1.6.1
tensorflow==2.4.1
